"""
Quick inference on trained Qwen2.5 VL model
"""

# from align_anything.models.qwen2_5_vl import AccustomedQwen2_5_VLModel
import json
from transformers import (
    AutoTokenizer,
    AutoProcessor,
)

try:
    from qwen_vl_utils import process_vision_info
except ImportError as e:
    pass
try:
    from transformers.models.qwen2_5_vl import Qwen2_5_VLModel
    from transformers import Qwen2_5_VLForConditionalGeneration
except ImportError as e:
    pass
from loguru import logger
from transformers import TextStreamer

import os
from loguru import logger

import torch
import glob
import natsort
import sys


class VLBase:
    def __init__(self, model_path=None, processor_path=None, device="auto"):
        self.device = (
            "cuda"
            if torch.cuda.is_available()
            else (
                "mps"
                if torch.backends.mps.is_available()
                else "cpu" if device == "auto" else device
            )
        )
        self.model = self.load_model(model_path)
        self.processor = self.load_processor(
            processor_path if processor_path is None else model_path
        )

        self.history_msgs = []

    def load_model(self, model_path):
        raise NotImplementedError

    def load_processor(self, processor_path):
        raise NotImplementedError

    def stream_chat_with_images(self):
        raise NotImplementedError

    def generate(self, prompt, image, verbose):
        pass


class Qwen2_5_VL(VLBase):
    def __init__(self, model_path=None, processor_path=None, device="auto"):
        super().__init__(model_path, processor_path, device)
        # default: Load the model on the available device(s)

    def load_model(self, model_path):
        if model_path is None:
            model_path = "checkpoints/Qwen2.5-VL-7B-Instruct"
        model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
            model_path,
            torch_dtype="bfloat16",
            attn_implementation="flash_attention_2",
        )
        model.to(self.device)
        logger.info(f"model loaded from: {model_path}")
        return model

    def load_processor(self, processor_path):
        if processor_path is None:
            processor_path = "checkpoints/Qwen2.5-VL-3B-Instruct"
        processor = AutoProcessor.from_pretrained(processor_path)
        self.tokenizer = AutoTokenizer.from_pretrained(processor_path)
        if self.tokenizer.pad_token_id is None:
            self.tokenizer.pad_token_id = self.tokenizer.encode(
                self.tokenizer.pad_token
            )
        return processor

    def get_msg(self, text, image=None, system_msg=None):
        if image is None:
            return {
                "role": "user",
                "content": [
                    {"type": "text", "text": text},
                ],
            }
        elif os.path.isdir(image):
            image = glob.glob(os.path.join(image, "*.png"))
            image = natsort.sorted(image)

        if isinstance(image, list):
            a = {
                "role": "user",
                "content": [
                    {"type": "text", "text": text},
                ],
            }
            a["content"].append(
                {
                    "type": "video",
                    "video": image,
                    "max_pixels": 360 * 420,
                }
            )
            # for im in image:
            #     a['content'].append( {
            #             "type": "image",
            #             "image": im,
            #         },)
            return a

        return {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image": image,
                    "max_pixels": 430 * 28 * 28,
                },
                {"type": "text", "text": text},
            ],
        }

    def generate(
        self,
        prompt,
        images,
        stream=True,
        max_size=700,
        verbose=False,
        prevent_more_image=True,
        system_msg=None,
    ):
        msg = self.get_msg(prompt, images)
        messages = [msg]
        if system_msg:
            messages.append({"role": "system", "content": system_msg})

        # Preparation for inference
        text = self.processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        image_inputs, video_inputs = process_vision_info(messages)
        print(image_inputs, video_inputs)
        inputs = self.processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        )
        inputs = inputs.to(self.device)

        # Inference: Generation of the output
        streamer = TextStreamer(
            self.tokenizer, skip_prompt=True, skip_special_tokens=True
        )

        generated_ids = self.model.generate(
            **inputs, do_sample=False, max_new_tokens=500, streamer=streamer
        )
        generated_ids_trimmed = [
            out_ids[len(in_ids) :]
            for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
        ]
        output_text = self.processor.batch_decode(
            generated_ids_trimmed,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=False,
        )
        # print(output_text)
        return output_text


def chat():
    from termcolor import colored

    model_p = sys.argv[1]
    model = Qwen2_5_VL(model_path=model_p, processor_path=model_p)

    crt_input = ["images/cats.jpg", None]

    while True:
        img_or_txt = input(colored("\nUser (txt/img_path): ", "cyan")).strip()

        if os.path.exists(img_or_txt.split(" ")[0]):
            crt_input[0] = img_or_txt
            print(colored("System: Image updated.", "green"))
            continue
        else:
            crt_input[1] = img_or_txt

        if crt_input[0] and crt_input[1]:
            print(colored("Assistant:", "green"), end=" ")
            model.generate(images=crt_input[0], prompt=crt_input[1], verbose=False)
            crt_input[0] = None
        elif not crt_input[0] and crt_input[1]:
            # pure text
            print(colored("Assistant:", "green"), end=" ")
            model.generate(images=None, prompt=crt_input[1], verbose=False)
        else:
            print(
                colored("System: Please provide either an image or text input.", "red")
            )


if __name__ == "__main__":
    # infer()
    chat()
